package tool

import (
	"fmt"
	"log"
	"os"

	"github.com/tuneinsight/lattigo/v4/bfv"
	"github.com/tuneinsight/lattigo/v4/rlwe"
)

// **生成 BFV 密钥对**
func GenerateBFVKeys() (bfv.Parameters, *rlwe.SecretKey, *rlwe.PublicKey) {
	// 1️⃣ 使用官方推荐参数
	// paramsLiteral := bfv.ParametersLiteral{
	// 	LogN: 13,                                            // N = 2^13 = 8192，影响计算复杂度
	// 	Q:    []uint64{0xffffffffffc0001, 0xfffffffff00001}, // 质数模数链
	// 	P:    []uint64{0x1ffffffffffffff},                   // ✅ 这里要选素数
	// 	T:    65537,                                         // 明文模数（素数）
	// }

	// 2️⃣ 通过 `NewParametersFromLiteral()` 创建参数
	params, err := bfv.NewParametersFromLiteral(bfv.PN13QP218)
	if err != nil {
		log.Fatalf("参数初始化失败: %v", err)
	}

	// 2️⃣ 直接使用 `bfv.NewKeyGenerator()`，而不是 `rlwe.NewKeyGenerator()`
	kgen := bfv.NewKeyGenerator(params)

	// 3️⃣ 生成密钥对
	sk, pk := kgen.GenKeyPair()

	return params, sk, pk
}

// **加密文件**
func EncryptFileBFV(filename string, params bfv.Parameters, pk *rlwe.PublicKey) []*rlwe.Ciphertext {
	// 读取文件内容
	data, err := os.ReadFile(filename)
	if err != nil {
		log.Fatalf("文件读取失败: %v", err)
	}

	// 将字节转换成整数
	intData := BytesToInts(data)

	// 创建编码器 & 加密器
	encoder := bfv.NewEncoder(params)
	encryptor := bfv.NewEncryptor(params, pk)

	// BFV 方案的可用槽位 (N/2)
	batchSize := params.N() / 2

	var ciphertexts []*rlwe.Ciphertext
	for i := 0; i < len(intData); i += batchSize {
		// 取 batchSize 大小的数据
		end := i + batchSize
		if end > len(intData) {
			end = len(intData)
		}

		// 创建明文 & 编码
		pt := bfv.NewPlaintext(params, params.MaxLevel())
		encoder.Encode(intData[i:end], pt)

		// 加密
		ct := encryptor.EncryptNew(pt)
		ciphertexts = append(ciphertexts, ct)
	}

	fmt.Println("文件加密完成！")
	return ciphertexts
}

// **字节转换整数**
func BytesToInts(data []byte) []int64 {
	result := make([]int64, len(data))
	for i, b := range data {
		result[i] = int64(b)
	}
	return result
}

// **解密文件**
func DecryptFileBFV(ciphertexts []*rlwe.Ciphertext, params bfv.Parameters, sk *rlwe.SecretKey) []byte {
	decryptor := bfv.NewDecryptor(params, sk)
	encoder := bfv.NewEncoder(params)

	var decryptedData []byte
	for _, ct := range ciphertexts {
		pt := decryptor.DecryptNew(ct)
		values := encoder.DecodeIntNew(pt)

		// 过滤掉无效的 0 字节
		for _, v := range values {
			if v != 0 { // ✅ 过滤掉 `null` 或 `0x00`
				decryptedData = append(decryptedData, byte(v))
			}
		}
	}

	fmt.Println("文件解密完成！")
	return decryptedData
}
